import torch
import torch.nn as nn
import math
# torch中变量封装函数Variable.
from torch.autograd import Variable

# 嵌入层
# 定义Embeddings类来实现文本嵌入层，这里s说明代表两个一模一样的嵌入层，他们共享参数
# 该类继承nn.Module，这样就有标准层的一些功能，这里我们也可以理解为一种模式，我们自己实现的所有层
# 作用：把词汇的“编号”转换为向量表示，在高纬度空间捕捉词汇间的关系
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        """
        :param d_model: 词嵌入的维度
        :param vocab: 词表的大小
        """
        # 使用super的方式指明继承nn.Module的初始化函数
        super(Embeddings, self).__init__()

        # 使用nn中原版的Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        """
        :param x: 因为Embedding层是首层，所以代表输入给型的文本通过词汇映射后的张量
        :return:
        """
        # 该层的前向传播逻辑，所有层中都会有此函数当传给该类的实例化对象参数时，
        # 将x传给self.1ut并与根号下self.d_mode1相乘作为结果返回
        # 乘以d model 用于缩放
        return self.lut(x) * math.sqrt(self.d_model)


if __name__ == '__main__':

    embedding = nn.Embedding(10,3)
    input = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
    output = embedding.forward(input)
    print(output)

    output = embedding(input)
    print(output)

    # padding 参数确保0号单词 向量的坐标就在向量 0
    embedding = nn.Embedding(10,3, padding_idx=0)


    d_model = 512
    vocab = 1000
    x = Variable(torch.LongTensor([[100,2,421,508],[491,998,1,221]]))
    my_embeddings = Embeddings(d_model,  vocab)
    embedded_x = my_embeddings(x)
    #
    print(embedded_x)
    print(embedded_x.shape)


